{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Earrbj1HNaVk"
   },
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-ai-edge/model-explorer/blob/main/example_colabs/quick_start.ipynb)\n",
    "\n",
    "# Google AI Edge Model Explorer\n",
    "A visualization tool that lets you analyze ML models and graphs, accelerating deployment to on-device targets. [Learn more](https://ai.google.dev/edge/model-explorer).\n",
    "\n",
    "**Key Features**\n",
    "\n",
    "* Visualize large models effortlessly\n",
    "* Find model conversion issues\n",
    "* Identify optimization targets\n",
    "* Easy to use intuitive UI\n",
    "\n",
    "Follow the [installation instructions](https://github.com/google-ai-edge/model-explorer/wiki/5.-Run-in-Colab-Notebook) to add it to your own Colab.\n",
    "\n",
    "Want to run Model Explorer locally? [Get Started here](https://github.com/google-ai-edge/model-explorer/wiki/1.-Installation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vCM_sxQsg2nj"
   },
   "source": [
    "# Download a copy of the EfficientDet TFLite model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SmUNdu2jhU1z"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import ssl\n",
    "import tempfile\n",
    "import urllib.request\n",
    "import torch\n",
    "import torchvision\n",
    "\n",
    "# Create an SSL context that does not verify certificates.\n",
    "ssl_context = ssl._create_unverified_context()\n",
    "tmp_path = tempfile.mkdtemp()\n",
    "model_path = os.path.join(tmp_path, \"model.tflite\")\n",
    "\n",
    "# Download the model\n",
    "try:\n",
    "  with urllib.request.urlopen(\n",
    "      \"https://storage.googleapis.com/tfweb/model-graph-vis-v2-test-models/efficientdet.tflite\",\n",
    "      context=ssl_context,\n",
    "  ) as response:\n",
    "    data = response.read()\n",
    "\n",
    "  with open(model_path, \"wb\") as file:\n",
    "    file.write(data)\n",
    "\n",
    "  print(\"Model downloaded successfully!\")\n",
    "\n",
    "except Exception as e:\n",
    "  print(f\"Failed to download model: {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nSHa2LHCg7Gz"
   },
   "source": [
    "# Install Model Explorer using pip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XqnjhEVqkSvU"
   },
   "outputs": [],
   "source": [
    "!pip install ai-edge-model-explorer\n",
    "\n",
    "# Faster installation by skipping deps that are included in colab runtime:\n",
    "# !pip install --no-deps ai-edge-model-explorer-adapter ai-edge-model-explorer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tNzvaoFHhBJR"
   },
   "source": [
    "# Visualize the downloaded EfficientDet model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_qycf3tbmP_S"
   },
   "outputs": [],
   "source": [
    "import model_explorer\n",
    "\n",
    "model_explorer.visualize(model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Q3aMYkE6HMPM"
   },
   "source": [
    "# Visualize a PyTorch model\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sZJkEgdsHSH9"
   },
   "outputs": [],
   "source": [
    "# Get mobilnet v2 pytorch model as an example.\n",
    "model = torchvision.models.mobilenet_v2().eval()\n",
    "inputs = (torch.rand([1, 3, 224, 224]),)\n",
    "ep = torch.export.export(model, inputs)\n",
    "\n",
    "# Visualize\n",
    "model_explorer.visualize_pytorch('mobilenet', exported_program=ep)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualize a JAX model (via StableHLO MLIR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "\n",
    "\n",
    "def my_function(x):\n",
    "  y = jnp.sin(x)\n",
    "  z = jnp.cos(x)\n",
    "  return y * z\n",
    "\n",
    "\n",
    "inputs = jnp.array(2.0)\n",
    "\n",
    "# JIT, lower the function, and get the textual representation\n",
    "#   - jax.jit(fn): Creates a JIT-compiled version of our function.\n",
    "#   - .lower(inputs): Traces the function with the dummy input to generate\n",
    "#     the low-level representation without actually executing it.\n",
    "#   - .as_text(debug_info=True): Converts the lowered representation (StableHLO)\n",
    "#     to a human-readable string. `debug_info=True` includes source-level\n",
    "#     location information, which is useful for visualization.\n",
    "stablehlo_mlir = jax.jit(my_function).lower(inputs).as_text(debug_info=True)\n",
    "\n",
    "# Write the MLIR to file\n",
    "file_path = '/content/stablehlo.mlir'\n",
    "with open(file_path, 'w') as f:\n",
    "  f.write(stablehlo_mlir)\n",
    "\n",
    "# Visualize\n",
    "model_explorer.visualize(file_path)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [
    "vCM_sxQsg2nj",
    "nSHa2LHCg7Gz",
    "tNzvaoFHhBJR",
    "Q3aMYkE6HMPM"
   ],
   "private_outputs": true,
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
